{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Implementing an Agentic RAG System with Open Source SLMs and SmolAgents\n",
        "This notebook is a companion of chapter 13 of the \"Domain-Specific Small Language Models\" [book](https://www.manning.com/books/domain-specific-small-language-models), author Guglielmo Iozzia, [Manning Publications](https://www.manning.com/), 2025.  \n",
        "The code in this notebook is an example implementation of an Agentic RAG (Retrieval Augmented Generation) system using only Small Language Models (SLMs), the HF's [SmolAgents](https://github.com/huggingface/smolagents) framework, [LangChain](https://github.com/langchain-ai/langchain) and [LanceDB](https://github.com/lancedb/lancedb). Hardware acceleration (GPU) is recommended.   \n",
        "More details about the code can be found in the related book's chapter."
      ],
      "metadata": {
        "id": "waBrKsUhblS2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the missing requirements in the Colab VM."
      ],
      "metadata": {
        "id": "KdnJ0Ulay9-o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install smolagents langchain lancedb langchain-community rank_bm25 pypdf langchain-huggingface ddgs"
      ],
      "metadata": {
        "id": "OEQxoWrJ56Pw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Upload a PDF document."
      ],
      "metadata": {
        "id": "ixbb2AVIdBNB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!curl https://arxiv.org/pdf/2502.12923 --output arxiv_250212923.pdf"
      ],
      "metadata": {
        "id": "Wzk9YONdCMt8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Extract the text from the uploaded document and chunk it."
      ],
      "metadata": {
        "id": "_x5eJBbLdH0b"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.document_loaders import PyPDFLoader\n",
        "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
        "\n",
        "loader = PyPDFLoader(\"/content/arxiv_250212923.pdf\")\n",
        "documents = loader.load()\n",
        "\n",
        "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)\n",
        "docs = text_splitter.split_documents(documents)\n",
        "\n",
        "print(len(docs))"
      ],
      "metadata": {
        "id": "7z6cpiIHCG4H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download the embedding model, [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)."
      ],
      "metadata": {
        "id": "S1g05c4OdPFe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain_huggingface import HuggingFaceEmbeddings\n",
        "\n",
        "embeddings = HuggingFaceEmbeddings(model_name=\"all-MiniLM-L6-v2\")"
      ],
      "metadata": {
        "id": "spj0LKYoztfy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create a local LanceDB database and the table where to store the embeddings. Then use the embedding model to transform the text chunks into embeddings and store them into the `document_embeddings` table."
      ],
      "metadata": {
        "id": "4dMHPED3rQNH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import lancedb\n",
        "import numpy as np\n",
        "import pyarrow as pa\n",
        "\n",
        "db = lancedb.connect(\"./lancedb\")\n",
        "\n",
        "# Define the schema for the table\n",
        "schema = pa.schema([\n",
        "    pa.field(\"embedding\", pa.list_(pa.float32(), list_size=384)), # Specify FixedSizeList for embedding\n",
        "    pa.field(\"text\", pa.string()),\n",
        "])\n",
        "\n",
        "table = db.create_table(\n",
        "    \"document_embeddings\",\n",
        "    schema=schema, # Pass the schema to create_table\n",
        "    data=[\n",
        "        {\n",
        "            \"embedding\": np.array(embeddings.embed_query(doc.page_content), dtype=np.float32).flatten().tolist(),\n",
        "            \"text\": doc.page_content,\n",
        "        }\n",
        "        for doc in docs\n",
        "    ],\n",
        ")"
      ],
      "metadata": {
        "id": "EO_O-jjNg55J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom SmolAgents tool to perform BM25 search on the knowledge base. BM25, or Best Matching 25, is a scoring algorithm used by search engines to evaluate how well a document matches a specific search query. It ranks documents based on factors like term frequency, document length, and the rarity of terms, making it effective for information retrieval tasks."
      ],
      "metadata": {
        "id": "wRWp9eS6rwLg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import List\n",
        "from langchain.docstore.document import Document\n",
        "from smolagents import Tool\n",
        "from langchain_community.retrievers import BM25Retriever\n",
        "\n",
        "class BM25SearchTool(Tool):\n",
        "    name = \"do_bm25_search_on_local_documents\"\n",
        "    description = \"Uses text search to retrieve the parts of the documentation that could be most relevant to answer a query.\"\n",
        "    inputs = {\n",
        "        \"query\": {\n",
        "            \"type\": \"string\",\n",
        "            \"description\": \"The search query string.\"\n",
        "        },\n",
        "    }\n",
        "    output_type = \"string\"\n",
        "\n",
        "    def __init__(self, docs: List[Document], **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self.docs = docs\n",
        "\n",
        "    def forward(self, query: str) -> str:\n",
        "        retriever = BM25Retriever.from_documents(self.docs, k=3)\n",
        "        docs = retriever.invoke(query)\n",
        "        return \"\\nRetrieved documents:\\n\" + \"\".join(\n",
        "            [f\"\\n\\n===== Document {i} =====\\n\" + doc.page_content for i, doc in enumerate(docs)]\n",
        "        )"
      ],
      "metadata": {
        "id": "josRAfMrzTGB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define a custom SmolAgents tool that performs semantic search on the `document_embeddings` table."
      ],
      "metadata": {
        "id": "BDpLpEASwod7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from smolagents import Tool\n",
        "import numpy as np\n",
        "\n",
        "class SemanticSearchTool(Tool):\n",
        "    name = \"semantic_search\"\n",
        "    description = \"Performs semantic search on a database of document embeddings.\"\n",
        "    inputs = {\n",
        "        \"query\": {\n",
        "            \"type\": \"string\",\n",
        "            \"description\": \"The search query string.\",\n",
        "        },\n",
        "        \"top_k\": {\n",
        "            \"type\": \"integer\",\n",
        "            \"description\": \"The number of top results to return.\",\n",
        "            \"default\": 1,\n",
        "            \"nullable\": True,\n",
        "        },\n",
        "    }\n",
        "    output_type = \"string\"\n",
        "\n",
        "    def __init__(self, table, embeddings, **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self.table = table\n",
        "        self.embeddings = embeddings\n",
        "\n",
        "    def forward(self, query: str, top_k: int = 1) -> str:\n",
        "        \"\"\"Performs semantic search on the document_embeddings table.\"\"\"\n",
        "        query_embedding = self.embeddings.embed_query(query)\n",
        "        query_embedding = np.array(query_embedding)\n",
        "        results = self.table.search(query_embedding, vector_column_name=\"embedding\").limit(top_k).to_df()\n",
        "        return results.to_string()\n"
      ],
      "metadata": {
        "id": "dxeaFJqA7xTm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Implement a custom SmolAgents Tool that does hybrid search (BM25 keyword search + vector search) on a LanceDB table. It uses LangChain's `EnsableRetriever` to combine the search results."
      ],
      "metadata": {
        "id": "yf8bU0tAw4ik"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.retrievers import EnsembleRetriever\n",
        "from langchain_community.retrievers import BM25Retriever\n",
        "from smolagents import Tool\n",
        "\n",
        "class HybridSearchTool(Tool):\n",
        "    name = \"hybrid_search\"\n",
        "    description = \"Performs a hybrid search (BM25 keyword + vector search) on a LanceDB table.\"\n",
        "    inputs = {\n",
        "        \"query\": {\n",
        "            \"type\": \"string\",\n",
        "            \"description\": \"The search query string.\",\n",
        "        },\n",
        "        \"top_k\": {\n",
        "            \"type\": \"integer\",\n",
        "            \"description\": \"The number of top results to return.\",\n",
        "            \"default\": 1,\n",
        "            \"nullable\": True,\n",
        "        },\n",
        "    }\n",
        "    output_type = \"string\"\n",
        "\n",
        "    def __init__(self, bm25_retriever, semantic_retriever, **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self.ensemble_retriever = EnsembleRetriever(\n",
        "            retrievers=[bm25_retriever, semantic_retriever],\n",
        "            weights=[0.4, 0.6])\n",
        "\n",
        "    def forward(self, query: str, top_k: int = 1) -> str:\n",
        "        docs = self.ensemble_retriever.get_relevant_documents(query)\n",
        "        # Process the results as needed\n",
        "        return \"\\nRetrieved documents:\\n\" + \"\".join(\n",
        "            [f\"\\n\\n===== Document {i} =====\\n\" + doc.page_content for i, doc in enumerate(docs)]\n",
        "        )\n"
      ],
      "metadata": {
        "id": "bH2m5hhNYUtZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Implement a LanceDB custom retriever class."
      ],
      "metadata": {
        "id": "o8pcaeK2xS8T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import List\n",
        "from langchain.schema import BaseRetriever, Document\n",
        "from pydantic import Field\n",
        "\n",
        "class LanceDBVectorSearch(BaseRetriever):\n",
        "    \"\"\"\n",
        "    A vector search retriever that uses LanceDB for storage and retrieval.\n",
        "\n",
        "    Attributes:\n",
        "        table (lancedb.LanceTable): The LanceDB table to search.\n",
        "        embeddings (HuggingFaceEmbeddings): The embeddings model to use.\n",
        "    \"\"\"\n",
        "\n",
        "    table: lancedb.table.LanceTable = Field(...) # Add Field for table\n",
        "    embeddings: HuggingFaceEmbeddings = Field(...)\n",
        "\n",
        "    def __init__(self, table: lancedb.table.LanceTable, embeddings: HuggingFaceEmbeddings):\n",
        "        \"\"\"\n",
        "        Initializes the LanceDBVectorSearch retriever.\n",
        "\n",
        "        Args:\n",
        "            table (lancedb.LanceTable): The LanceDB table to search.\n",
        "            embeddings (HuggingFaceEmbeddings): The embeddings model to use.\n",
        "        \"\"\"\n",
        "        super().__init__(table=table, embeddings=embeddings) # Pass table and embeddings to super().__init__\n",
        "\n",
        "    def _get_relevant_documents(self, query: str) -> List[Document]:\n",
        "        \"\"\"\n",
        "        Retrieves relevant documents based on the given query.\n",
        "\n",
        "        Args:\n",
        "            query (str): The search query.\n",
        "\n",
        "        Returns:\n",
        "            List[Document]: A list of relevant documents.\n",
        "        \"\"\"\n",
        "        query_embedding = self.embeddings.embed_query(query)\n",
        "        query_embedding = np.array(query_embedding)\n",
        "\n",
        "        # Perform the search in LanceDB\n",
        "        results = self.table.search(query_embedding, vector_column_name=\"embedding\").limit(2).to_pandas()  # Limit to top 10 results\n",
        "\n",
        "        # Convert the results to Document objects\n",
        "        documents = [\n",
        "            Document(page_content=row[\"text\"], metadata={})\n",
        "            for _, row in results.iterrows()\n",
        "        ]\n",
        "\n",
        "        return documents"
      ],
      "metadata": {
        "id": "Pgt1JmW_dViX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Test the hybrid search."
      ],
      "metadata": {
        "id": "smWeM8E6x5TY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "bm25_retriever = BM25Retriever.from_documents(docs, k=3)\n",
        "semantic_retriever = LanceDBVectorSearch(table=table, embeddings=embeddings)\n",
        "\n",
        "hybrid_search_tool = HybridSearchTool(bm25_retriever, semantic_retriever)\n",
        "results = hybrid_search_tool.forward(\"Smart Home technologies\", top_k=2)\n",
        "print(results)"
      ],
      "metadata": {
        "id": "UP8XxDBXbFsE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Download the instructed SLM ([SmolLM2-1.7B-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct)) for the agent."
      ],
      "metadata": {
        "id": "D8Yzsskhx_x9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import torch\n",
        "from smolagents import TransformersModel, CodeAgent, DuckDuckGoSearchTool\n",
        "\n",
        "model_id = \"HuggingFaceTB/SmolLM2-1.7B-Instruct\"\n",
        "model = TransformersModel(model_id,\n",
        "                          device_map=\"auto\",\n",
        "                          max_new_tokens=700,\n",
        "                          torch_dtype=torch.float16)"
      ],
      "metadata": {
        "id": "lwEIjGLK1hYh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create an instance of the SmolAgents' built-in `DuckDuckGoSearchTool`."
      ],
      "metadata": {
        "id": "0yo-NCKq0u-u"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from smolagents import DuckDuckGoSearchTool\n",
        "\n",
        "duckduckgo_search_tool = DuckDuckGoSearchTool()"
      ],
      "metadata": {
        "id": "SA9Pug-QTE9m"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create and setup the code agent, specifying the SLM to use and the list of allowed custom tools."
      ],
      "metadata": {
        "id": "aSukDOcS07DG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "custom_tools = [hybrid_search_tool, duckduckgo_search_tool]\n",
        "agent = CodeAgent(\n",
        "    tools=custom_tools,\n",
        "    model=model,\n",
        "    max_steps=3,\n",
        "    verbosity_level=2,\n",
        "    add_base_tools=False\n",
        ")"
      ],
      "metadata": {
        "id": "cAijsc-Anxkz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Perform a query with the agent."
      ],
      "metadata": {
        "id": "SYGWKeh61Fat"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#query = \"Search the web for MSD Ireland\"\n",
        "#query = \"Do an hybrid search about Smart Home technologies\"\n",
        "query = \"Do an hybrid search about Smart Home technologies and then search the web about the same\"\n",
        "#query = \"Do a semantic search for MSD\"\n",
        "agent_output = agent.run(query)\n",
        "print(\"Final output:\")\n",
        "agent_output\n"
      ],
      "metadata": {
        "id": "2rSUkuf813j4"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}